package layers

import (
	act "lenet/network/activations"
	"math/rand"
)

type fullLayer struct {
	activate act.ActivationInterface
	lRate    float64
	inpDim   uint
	outDim   uint
	inputs   []float64
	outputs  []float64
	weights  [][]float64
}

func NewFullLayer(inpDim, outDim uint, lRate float64, activation act.ActivationInterface) *fullLayer {
	w := make([][]float64, outDim)
	for i := 0; i < int(outDim); i++ {
		w[i] = make([]float64, inpDim)
		for j := 0; j < int(inpDim); j++ {
			w[i][j] = rand.Float64() / 1000
		}
	}

	return &fullLayer{
		activate: activation,
		lRate:    lRate,
		inpDim:   inpDim,
		outDim:   outDim,
		weights:  w,
	}
}

func (l *fullLayer) Forward(input []float64) []float64 {
	//log.Println("Forward full")
	l.inputs = input
	l.outputs = make([]float64, 0)
	ret := make([]float64, 0)
	for _, weight := range l.weights {
		dot := dotProduct(weight, l.inputs)
		l.outputs = append(l.outputs, dot)
		ret = append(ret, l.activate.Activate(dot))
	}

	//log.Println("Forward full end")
	return ret
}

func (l *fullLayer) BackProp(err []float64) []float64 {
	//log.Println("Backprop full")
	nextErr := l.backPropogation(err)
	//log.Println("Backprop full mid")
	l.changeWeights(err)
	//log.Println("Backprop full end, ret len ", len(nextErr))
	return nextErr
}

func dotProduct(a, b []float64) float64 {
	var sum float64
	for i := 0; i < len(a); i++ {
		sum += a[i] * b[i]
	}
	return sum
}

func (l *fullLayer) backPropogation(err []float64) []float64 {
	newErr := make([]float64, 0)
	//log.Println(l.inpDim, l.outDim, len(l.weights), len(l.weights[0]), len(err))
	for i := 0; i < int(l.inpDim); i++ {
		tmp := 0.
		for j := 0; j < int(l.outDim); j++ {
			tmp += err[j] * l.weights[j][i]
		}
		newErr = append(newErr, tmp)
	}
	//log.Println("newErr len:", len(newErr))
	return newErr
}

func (l *fullLayer) changeWeights(err []float64) {
	for n, arrNeuronWeights := range l.weights {
		for w := range arrNeuronWeights {
			arrNeuronWeights[w] += l.lRate * err[n] * l.inputs[w] * l.activate.Activate(l.outputs[n])
		}
	}
}
